## I add some code into simulated0529.R to generate correlation between factor loadings and network ties.
## In this function, corH controls the degree of correlation
library(MASS)
spatial.simulate <- function(W, ## W matrix
                             rhoZ,
	                         beta,
                             rhoA,
                             p = 2, 
	                         r = 0,
                             contextual.effect,
                             force,
	                         sigma2,
	                         kappa,  
	                         sigma_n2, corH) {

	library(abind)
    
    TT <- dim(W)[3]
	N <- dim(W)[1]

	## gen rho 
	Lambda <- matrix(NA, 2, TT)
	for (i in 1:TT) {
        sublam <- eigen(W[,,i])[[1]]
        pos <- which(Im(sublam) == 0) 
        
        sublam <- Re(sublam[pos])

        Lambda[1, i] <- max(sublam)
		Lambda[2, i] <- min(min(sublam), -1)
	}

	thres <- matrix(NA, TT, 2)
	#thres[,1] <- 1 / apply(Lambda, 2, min)
	#thres[,2] <- 1 / apply(Lambda, 2, max)
    thres[,1] <- -1
    thres[,2] <- 1

	Rho <- matrix(0, TT, 1)

    rhoMu <- rhoZ %*% rhoA

    Rhot <- NULL
    con1 <- FALSE

    while(!con1) {
    	
        if (kappa == 0) {
            Rhot <- rnorm(TT, sd = sqrt(sigma_n2)) + c(rhoMu)
        } else {
            
            Rhot <- rep(NA, TT)
            Rhot[1] <- rhoMu[1] + rnorm(1, sd = sqrt(sigma_n2))
            for (i in 2:TT) {
                Rhot[i] <- rhoMu[i] + kappa * Rhot[i-1] + rnorm(1, sd = sqrt(sigma_n2))
            }
        }

    	a <- 0
    	for (i in 1:TT) {
    		rhot <- Rhot[i]
    		if (thres[i, 1] < rhot && rhot < thres[i, 2]) {
    			a <- a + 1
    		}
    	}

    	if (a == TT) {
    		con1 <- TRUE
    	}
    }

    Rho[,1] <- Rhot

    I <- diag(N)

    AW <- array(NA, dim = c(N, N, TT))
    for (i in 1:TT) {
    	AW[,,i] <- solve(I - (Rho[i,1]) * W[,,i])
    }


    p.old <- p
    
    ## gen X 
    X <- NULL
    if (p > 0) {
        X <- array(rnorm(N*p*TT), dim = c(N, p, TT))
    } 

    ## contextual 
    X2 <- X3 <- X1 <- NULL
    if (p > 0) {
        ## over all
        X1 <- array(NA, dim = c(N, p, TT))
        for (i in 1:TT) {
            X1[,,i] <- W[,,i] %*% X[,,i]
        }
        ## unit 
        X2 <- array(NA, dim = c(TT, p, N))
        for (i in 1:p) {
            X2[,i,] <- t(X1[,i,])
        }

        ## time
        X3 <- X1

    }

    xfit <- zfit <- afit <- matrix(0, N, TT)


    X.all <- NULL
    if (p == 0) {
        X.all <- array(1, dim = c(N, 1, TT))
    } else {
        if (contextual.effect == "none") {
            X.all <- abind(array(1, dim = c(N, 1, TT)), X, along = 2)
        } else {
            X.all <- abind(array(1, dim = c(N, 1, TT)), X, X1, along = 2)
        }
    }


    Z.all <- NULL
    if (p == 0) {
        if (force %in% c("unit", "both")) {
            Z.all <- array(1, dim = c(TT, 1, N))
        }
    } else {
        if (contextual.effect %in% c("unit", "both")) {
            if (force %in% c("unit", "both")) {
                Z.all <- abind(array(1, dim = c(TT, 1, N)), X2, along = 2)
            } else {
                Z.all <- X2
            }
        } 
        #else {
        #    if (force %in% c("unit", "both")) {
        #        Z.all <- array(1, dim = c(TT, 1, N))
        #    }
        #}
    }

    A.all <- NULL
    if (p == 0) {
        if (force %in% c("time", "both")) {
            A.all <- array(1, dim = c(N, 1, TT))
        }
    } else {
        if (contextual.effect %in% c("time", "both")) {
            if (force %in% c("time", "both")) {
                A.all <- abind(array(1, dim = c(N, 1, TT)), X3, along = 2)
            } else {
                A.all <- X3
            }
        } 
        #else {
        #    if (force %in% c("time", "both")) {
        #        A.all <- array(1, dim = c(N, 1, TT))
        #    }
        #}
    }


    if (!is.null(X.all)) {
    	p <- dim(X.all)[2]
        if (dim(beta)[1] != p) {
            stop(paste("Beta should be of length:", p, "Did you add the contextual effect?", sep = " "))
        }
        for (i in 1:TT) {
    		subx <- X.all[,,i]
    		if (p == 1) {
    			subx <- matrix(subx, N, 1)
    		}
    		xfit[,i] <- c(subx %*% beta)
    	}
    }

    if (!is.null(Z.all)) {
    	p <- dim(Z.all)[2]
        alpha <- matrix(rnorm(N*p), p, N)
    	for (i in 1:N) {
    		subz <- Z.all[,,i]
    		subal <- as.matrix(alpha[, i])
    		if (p == 1) {
    			subz <- matrix(subz, TT, 1)
    		}
    		zfit[i,] <- c(subz %*% subal)
    	}
    }

    if (!is.null(A.all)) {
    	p <- dim(A.all)[2]
        xi <- matrix(rnorm(TT*p), p, TT)
    	for (i in 1:TT) {
    		suba <- A.all[,,i]
    		subxi <- as.matrix(xi[, i])
    		if (p == 1) {
    			suba <- matrix(suba, N, 1)
    		}
    		afit[,i] <- c(suba %*% subxi)
    	}
    }

    yfit <- xfit + zfit + afit

    
       ## Ng <- max(group)
    fefit <- matrix(0, N, TT)
    mW <- W[ , ,1]
    for (i in 2: TT){
        mW <- mW+ W[, , i]
    }
    mW <- mW/TT
    mW[mW>0] <- corH
    diag(mW) <- 1
    mL <- rep(0, N)
    if (r > 0) {
        if(corH==0){
                    F <- matrix(rnorm(TT * r), TT, r)
                    L <- matrix(rnorm(N * r), N, r)
        }else{
                    F <- matrix(rnorm(TT * r), TT, r)
                    L <- t(mvrnorm(r, mL, mW))
        }
        fefit <- L %*% t(F)
        yfit <- yfit + fefit
    }

    e <- matrix(rnorm(TT*N, sd = sqrt(sigma2)), N, TT)

    Y <- yfit + e

    for (i in 1:TT) {
    	Y[,i] <- c(AW[,,i] %*% as.matrix(Y[,i]))
    }

    x <- NULL
    if (p.old > 0) {
        x <- matrix(NA, TT * N, p.old)
        for (i in 1:p.old) {
            x[, i] <- c(X[,i,])
        }
    }

    data <- cbind.data.frame(rep(1:N, TT), rep(1:TT, each = N), c(Y), c(e), x)
    data.name <- c("id", "time", "y", "e")
    if (p.old > 0) {
        for (i in 1:p.old) {
            data.name <- c(data.name, paste("x", i, sep = ""))
        }  
    }

    names(data) <- data.name

    data <- data[order(data[, "id"], data[, "time"]), ]

    return(list(data = data, Rho = Rho))
}




### Rho is equal to absolute zero
spatial.simulate2 <- function(W, ## W matrix
                             rhoZ,
	                         beta,
                             p = 2, 
	                         r = 0,
                             contextual.effect,
                             force,
	                         sigma2,
	                         kappa,  
	                         sigma_n2, corH) {

	library(abind)
    
    TT <- dim(W)[3]
	N <- dim(W)[1]

    ## gen rho
    
	Lambda <- matrix(NA, 2, TT)
	for (i in 1:TT) {
        sublam <- eigen(W[,,i])[[1]]
        pos <- which(Im(sublam) == 0) 
        
        sublam <- Re(sublam[pos])

        Lambda[1, i] <- max(sublam)
		Lambda[2, i] <- min(min(sublam), -1)
	}

	thres <- matrix(NA, TT, 2)
	#thres[,1] <- 1 / apply(Lambda, 2, min)
	#thres[,2] <- 1 / apply(Lambda, 2, max)
    thres[,1] <- -1
    thres[,2] <- 1

	Rho <- matrix(kappa, TT, 1)
      

    I <- diag(N)

    AW <- array(NA, dim = c(N, N, TT))
    for (i in 1:TT) {
    	AW[,,i] <- solve(I - (Rho[i,1]) * W[,,i])
    }


    p.old <- p
    
    ## gen X 
    X <- NULL
    if (p > 0) {
        X <- array(rnorm(N*p*TT), dim = c(N, p, TT))
    } 

    ## contextual 
    X2 <- X3 <- X1 <- NULL
    if (p > 0) {
        ## over all
        X1 <- array(NA, dim = c(N, p, TT))
        for (i in 1:TT) {
            X1[,,i] <- W[,,i] %*% X[,,i]
        }
        ## unit 
        X2 <- array(NA, dim = c(TT, p, N))
        for (i in 1:p) {
            X2[,i,] <- t(X1[,i,])
        }

        ## time
        X3 <- X1

    }

    xfit <- zfit <- afit <- matrix(0, N, TT)


    X.all <- NULL
    if (p == 0) {
        X.all <- array(1, dim = c(N, 1, TT))
    } else {
        if (contextual.effect == "none") {
            X.all <- abind(array(1, dim = c(N, 1, TT)), X, along = 2)
        } else {
            X.all <- abind(array(1, dim = c(N, 1, TT)), X, X1, along = 2)
        }
    }


    Z.all <- NULL
    if (p == 0) {
        if (force %in% c("unit", "both")) {
            Z.all <- array(1, dim = c(TT, 1, N))
        }
    } else {
        if (contextual.effect %in% c("unit", "both")) {
            if (force %in% c("unit", "both")) {
                Z.all <- abind(array(1, dim = c(TT, 1, N)), X2, along = 2)
            } else {
                Z.all <- X2
            }
        } 
        #else {
        #    if (force %in% c("unit", "both")) {
        #        Z.all <- array(1, dim = c(TT, 1, N))
        #    }
        #}
    }

    A.all <- NULL
    if (p == 0) {
        if (force %in% c("time", "both")) {
            A.all <- array(1, dim = c(N, 1, TT))
        }
    } else {
        if (contextual.effect %in% c("time", "both")) {
            if (force %in% c("time", "both")) {
                A.all <- abind(array(1, dim = c(N, 1, TT)), X3, along = 2)
            } else {
                A.all <- X3
            }
        } 
        #else {
        #    if (force %in% c("time", "both")) {
        #        A.all <- array(1, dim = c(N, 1, TT))
        #    }
        #}
    }


    if (!is.null(X.all)) {
    	p <- dim(X.all)[2]
        if (dim(beta)[1] != p) {
            stop(paste("Beta should be of length:", p, "Did you add the contextual effect?", sep = " "))
        }
        for (i in 1:TT) {
    		subx <- X.all[,,i]
    		if (p == 1) {
    			subx <- matrix(subx, N, 1)
    		}
    		xfit[,i] <- c(subx %*% beta)
    	}
    }

    if (!is.null(Z.all)) {
    	p <- dim(Z.all)[2]
        alpha <- matrix(rnorm(N*p), p, N)
    	for (i in 1:N) {
    		subz <- Z.all[,,i]
    		subal <- as.matrix(alpha[, i])
    		if (p == 1) {
    			subz <- matrix(subz, TT, 1)
    		}
    		zfit[i,] <- c(subz %*% subal)
    	}
    }

    if (!is.null(A.all)) {
    	p <- dim(A.all)[2]
        xi <- matrix(rnorm(TT*p), p, TT)
    	for (i in 1:TT) {
    		suba <- A.all[,,i]
    		subxi <- as.matrix(xi[, i])
    		if (p == 1) {
    			suba <- matrix(suba, N, 1)
    		}
    		afit[,i] <- c(suba %*% subxi)
    	}
    }

    yfit <- xfit + zfit + afit

    
       ## Ng <- max(group)
    fefit <- matrix(0, N, TT)
    mW <- W[ , ,1]
    for (i in 2: TT){
        mW <- mW+ W[, , i]
    }
    mW <- mW/TT
    mW[mW>0] <- corH
    diag(mW) <- 1
    mL <- rep(0, N)
    if (r > 0) {
        if(corH==0){
                    F <- matrix(rnorm(TT * r), TT, r)
                    L <- matrix(rnorm(N * r), N, r)
        }else{
                    F <- matrix(rnorm(TT * r), TT, r)
                    L <- t(mvrnorm(r, mL, mW))
        }
        fefit <- L %*% t(F)
        yfit <- yfit + fefit
    }

    e <- matrix(rnorm(TT*N, sd = sqrt(sigma2)), N, TT)

    Y <- yfit + e

    for (i in 1:TT) {
    	Y[,i] <- c(AW[,,i] %*% as.matrix(Y[,i]))
    }

    x <- NULL
    if (p.old > 0) {
        x <- matrix(NA, TT * N, p.old)
        for (i in 1:p.old) {
            x[, i] <- c(X[,i,])
        }
    }

    data <- cbind.data.frame(rep(1:N, TT), rep(1:TT, each = N), c(Y), c(e), x)
    data.name <- c("id", "time", "y", "e")
    if (p.old > 0) {
        for (i in 1:p.old) {
            data.name <- c(data.name, paste("x", i, sep = ""))
        }  
    }

    names(data) <- data.name

    data <- data[order(data[, "id"], data[, "time"]), ]

    return(list(data = data, Rho = Rho))
}
